package net.cyue.ort.util;

import com.sun.jna.Platform;
import com.sun.jna.platform.win32.Advapi32Util;
import com.sun.jna.platform.win32.WinReg;

public class ONNXRuntimeBackendUtil {
    public static boolean isSupportDirectML() {
        if (!Platform.isWindows()) {
            return false;
        }
        String key = "SOFTWARE\\Microsoft\\DirectML";
        return Advapi32Util.registryKeyExists(WinReg.HKEY_LOCAL_MACHINE, key);
    }

    public static boolean isSupportCUDA() {
        try {
            // 获取系统属性
            String osName = System.getProperty("os.name").toLowerCase();

            if (osName.contains("windows")) {
                // Windows系统下检查nvcc是否存在
                Process process = Runtime.getRuntime().exec("where nvcc");
                return process.waitFor() == 0;
            } else {
                // Linux/Mac系统下检查nvidia-smi命令
                Process process = Runtime.getRuntime().exec("nvidia-smi");
                return process.waitFor() == 0;
            }
        } catch (Exception e) {
            return false;
        }
    }
}
